package org.example.ndarray.operation;

import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.Shape;

public class Example2 {
    public static void main(String[] args) {
        try (NDManager ndManager = NDManager.newBaseManager()) {
            /*
              我们可以将多个NDArrays concatenate联结起来，将它们首位相接地堆叠起来，生成一个更大的NDArray
              只需要提供NDArray的列表然后告诉系统沿着哪个轴进行联结
              示例展示了当我们将两个矩阵分别沿0轴（行）和1轴（列）联结将会发生什么
             */
            NDArray x = ndManager.arange(12f).reshape(3, 4);
            NDArray y = ndManager.create(new float[]{2, 1, 4, 3, 1, 2, 3, 4, 4, 3, 2, 1}, new Shape(3, 4));
            System.out.println("x = " + x);
            System.out.println("y = " + y);
            // 默认 axis = 0
            System.out.println(x.concat(y));
            System.out.println(x.concat(y, 1));

            /*
              有时，我们想通过逻辑声明构建一个 NDArray
              例如，按照x与y对应位置是否相等生成一个新 NDArray
              又或是对 NDArray 中的所有元素求和，生成一个只含有一个元素的 NDArray
             */

            // 按照x与y对应位置是否相等生成一个新 NDArray
            System.out.println(x.eq(y));

            // 对 x 中的所有元素求和
            System.out.println(x.sum());
        }
    }
}
